package server

import (
	"sync"
	"time"
)

// ClientInfo 存储客户端信息
type ClientInfo struct {
	ClientAddr    string    `json:"clientAddr"`    // 客户端地址
	TargetAddr    string    `json:"targetAddr"`    // 目标服务器地址
	ConnectTime   time.Time `json:"connectTime"`   // 连接建立时间
	BytesSent     int64     `json:"bytesSent"`     // 发送字节数
	BytesReceived int64     `json:"bytesReceived"` // 接收字节数
	Status        int       `json:"status"`        // 连接状态
}

// ClientManager 管理所有连接的客户端
type ClientManager struct {
	clients map[string]*ClientInfo
	mu      sync.RWMutex
}

// NewClientManager 创建一个新的客户端管理器
func NewClientManager() *ClientManager {
	return &ClientManager{
		clients: make(map[string]*ClientInfo),
	}
}

// AddClient 添加客户端信息
func (cm *ClientManager) AddClient(info *ClientInfo) {
	cm.mu.Lock()
	defer cm.mu.Unlock()
	cm.clients[info.ClientAddr] = info
}

// RemoveClient 删除客户端信息
func (cm *ClientManager) RemoveClient(clientAddr string) {
	cm.mu.Lock()
	defer cm.mu.Unlock()
	delete(cm.clients, clientAddr)
}

// UpdateClientStats 更新客户端信息
func (cm *ClientManager) UpdateClientStats(clientAddr string, targetAddr string, status int) {
	cm.mu.Lock()
	defer cm.mu.Unlock()
	if client, ok := cm.clients[clientAddr]; ok {
		client.TargetAddr = targetAddr
		client.Status = status
	}
}

// UpdateClientBytesSent 更新客户端上传流量
func (cm *ClientManager) UpdateClientBytesSent(clientAddr string, bytesSent int64) {
	cm.mu.Lock()
	defer cm.mu.Unlock()
	if client, ok := cm.clients[clientAddr]; ok {
		client.BytesSent += bytesSent
	}
}

// UpdateClientBytesReceived 更新客户端下载流量
func (cm *ClientManager) UpdateClientBytesReceived(clientAddr string, bytesReceived int64) {
	cm.mu.Lock()
	defer cm.mu.Unlock()
	if client, ok := cm.clients[clientAddr]; ok {
		client.BytesReceived += bytesReceived
	}
}

// GetClient 获取单个客户端信息
func (cm *ClientManager) GetClient(id string) (*ClientInfo, bool) {
	cm.mu.RLock()
	defer cm.mu.RUnlock()
	client, ok := cm.clients[id]
	return client, ok
}

// GetAllClients 获取所有客户端信息
func (cm *ClientManager) GetAllClients() []*ClientInfo {
	cm.mu.RLock()
	defer cm.mu.RUnlock()
	clients := make([]*ClientInfo, 0, len(cm.clients))
	for _, client := range cm.clients {
		clients = append(clients, client)
	}
	return clients
}
